﻿
#define SPADAS_DEPRECATED_HIDE
#include "NextGPUEncoder.h"
#include "NextGPUDecoder.h"
#include "ffmpegdl.h"

namespace videonext
{
	const UInt OPERATION_WAIT_TIME = 1;

	class GPUEncoderVars : public Vars
	{
	public:
		Pointer encoder;
		Int sentCount;
		Stream<Double> times;
		Size2D frameSize;
		Pointer packet;
		Bool operationWait;
		Bool useHardwareFrame;
		Array<Pointer> hwFrames;
		UInt curHwFrame;
	};
}

using namespace videonext;

GPUEncoder::GPUEncoder() : Object<GPUEncoderVars>(new GPUEncoderVars, TRUE)
{
	vars->encoder = 0;
	vars->packet = 0;
	vars->curHwFrame = 0;
}
bool GPUEncoder::open(VideoDataCodec codec, AVHWDeviceType type, Size2D size, UInt bitRate)
{
	close();

	if (getFFMPEGLibraryVersion() < 3) return FALSE;

	String codecName;
	switch (codec)
	{
	case VideoDataCodec::H264:
		codecName = "h264";
		break;
	case VideoDataCodec::H265:
		codecName = "hevc";
		break;
	case VideoDataCodec::MJPEG:
		codecName = "mjpeg";
		break;
	default:
		return FALSE;
	}

	String hwName;
	switch (type)
	{
	case AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA:
		hwName = "_nvenc";
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_VIDEOTOOLBOX:
		hwName = "_videotoolbox";
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_QSV:
		hwName = "_qsv";
		break;
	case (AVHWDeviceType)9: // AV_HWDEVICE_TYPE_AMF
		hwName = "_amf";
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI:
		hwName = "_vaapi";
		break;
	default:
		return FALSE;
	}
	codecName += hwName;

	Int maxBFrames = 1;
	AVPixelFormat pixelFormat = AVPixelFormat::AV_PIX_FMT_NV12;
	vars->useHardwareFrame = FALSE;
	vars->operationWait = FALSE;
	if (codecName == "h264_qsv")
	{
		if (bitRate == 0) bitRate = size.width * size.height * 8;
	}
	else if (codecName == "hevc_qsv")
	{
		if (bitRate == 0) bitRate = size.width * size.height * 4;
	}
	else if (codecName == "h264_nvenc")
	{
		maxBFrames = 0;
	}
	else if (codecName == "hevc_nvenc")
	{
		maxBFrames = 0;
	}
	else if (codecName == "h264_amf")
	{
		maxBFrames = 0;
		vars->operationWait = TRUE;
	}
	else if (codecName == "hevc_amf")
	{
		maxBFrames = 0;
		vars->operationWait = TRUE;
	}
	else if (codecName.endsWith("_vaapi"))
	{
		pixelFormat = AVPixelFormat::AV_PIX_FMT_VAAPI;
		vars->useHardwareFrame = TRUE;
	}

	Pointer targetCodec = avcodec_find_encoder_by_name(codecName.chars().data());
	if (targetCodec == 0) return FALSE;

	Pointer encoderTmp = avcodec_alloc_context3(targetCodec);
	if (encoderTmp == 0) return FALSE;

	const UInt PoolSize = 32;
	Pointer hwFrameCtx = NULL;
	if (vars->useHardwareFrame)
	{
		Pointer hardwareTmp = NULL;
		GPUDecoder::getHardwareCreateLocks()[(UInt)type].enter();
		Int resultCreateHardware = av_hwdevice_ctx_create(&hardwareTmp, type, NULL, NULL, 0);
		GPUDecoder::getHardwareCreateLocks()[(UInt)type].leave();
		if (resultCreateHardware < 0)
		{
			avcodec_free_context(&encoderTmp);
			return FALSE;
		}

		#if FFMPEGDL_VER_MAJOR > 1 || FFMPEGDL_VER_MINOR >= 7
		hwFrameCtx = av_hwframe_ctx_alloc(hardwareTmp);
		if (!hwFrameCtx)
		{
			avcodec_free_context(&encoderTmp);
			return FALSE;
		}

		Pointer hwFrameCtxData = av_hwframe_ctx_get_data(hwFrameCtx);
		av_hwframe_ctx_data_set_format(hwFrameCtxData, pixelFormat);
		av_hwframe_ctx_data_set_sw_format(hwFrameCtxData, AVPixelFormat::AV_PIX_FMT_NV12);
		av_hwframe_ctx_data_set_width_align32(hwFrameCtxData, size.width);
		av_hwframe_ctx_data_set_height_align32(hwFrameCtxData, size.height);
		av_hwframe_ctx_data_set_initial_pool_size(hwFrameCtxData, PoolSize);

		Int ret = av_hwframe_ctx_init(hwFrameCtx);
		if (ret >= 0)
		{
			avcodec_set_hw_frames_ctx(encoderTmp, hwFrameCtx);
			avcodec_set_hw_device_ctx(encoderTmp, hardwareTmp);
		}
		else
		{
			av_buffer_unref(&hwFrameCtx);
			avcodec_free_context(&encoderTmp);
			return FALSE;
		}
		#else
		avcodec_free_context(&encoderTmp);
		return FALSE;
		#endif
	}

	avcodec_set_bit_rate(encoderTmp, bitRate);
	avcodec_set_width(encoderTmp, size.width);
	avcodec_set_height(encoderTmp, size.height);
	avcodec_set_time_base(encoderTmp, 1, 30);
	avcodec_set_framerate(encoderTmp, 30, 1);
	avcodec_set_gop_size(encoderTmp, 32);
	avcodec_set_max_b_frames(encoderTmp, maxBFrames);
	avcodec_set_pix_fmt(encoderTmp, pixelFormat);

	Int resultOpenCodec = avcodec_open2(encoderTmp, targetCodec, NULL);
	if (resultOpenCodec < 0)
	{
		avcodec_free_context(&encoderTmp);
		return FALSE;
	}

	if (vars->useHardwareFrame)
	{
		#if FFMPEGDL_VER_MAJOR > 1 || FFMPEGDL_VER_MINOR >= 11
		vars->hwFrames = Array<Pointer>(PoolSize);
		for (UInt i = 0; i < PoolSize; i++)
		{
			vars->hwFrames[i] = av_frame_alloc();
			av_hwframe_get_buffer(hwFrameCtx, vars->hwFrames[i], 0);
		}
		vars->curHwFrame = 0;
		#else
		avcodec_free_context(&encoderTmp);
		return FALSE;
		#endif
	}

	vars->packet = av_packet_alloc();
	if (!vars->packet)
	{
		avcodec_free_context(&encoderTmp);
		return FALSE;
	}

	vars->encoder = encoderTmp;
	vars->sentCount = 0;
	vars->times = Stream<Double>(8);
	vars->frameSize = size;

	return TRUE;
}
void GPUEncoder::close()
{
	if (vars->encoder)
	{
		avcodec_free_context(&vars->encoder);
		for (auto e = vars->hwFrames.firstElem(); e.valid(); ++e) av_frame_free(e.value());
		vars->hwFrames = Array<Pointer>();
		vars->curHwFrame = 0;
	}
	if (vars->packet)
	{
		av_packet_free(&vars->packet);
	}
}
Bool GPUEncoder::send(CompatibleImage image, Double time)
{
	if (!vars->encoder) return FALSE;
	if (image.width() != vars->frameSize.width || image.height() != vars->frameSize.height) return FALSE;
	if (image.getType() != CompatibleImageType::AVFrameNV12) return FALSE;

	Pointer frame = image.getImage();
	if (vars->useHardwareFrame)
	{
		Pointer hwFrame = vars->hwFrames[vars->curHwFrame++];
		if (vars->curHwFrame >= vars->hwFrames.size()) vars->curHwFrame = 0;

		Int resultTransferData = av_hwframe_transfer_data(hwFrame, frame, 0);
		if (resultTransferData < 0)
		{
			av_frame_free(hwFrame);
			return FALSE;
		}
		frame = hwFrame;
	}

	Int resultSendFrame = avcodec_send_frame(vars->encoder, frame);
	if (resultSendFrame >= 0)
	{
		vars->sentCount++;
		vars->times.enqueue(time);
		if (vars->operationWait) Flag(OPERATION_WAIT_TIME).waitSet();
		return TRUE;
	}
	return FALSE;
}
void GPUEncoder::flush()
{
	if (vars->encoder == 0) return;
	avcodec_send_frame(vars->encoder, NULL);
	if (vars->operationWait) Flag(OPERATION_WAIT_TIME).waitSet();
}
Bool GPUEncoder::receive(Binary& packetData, Double& time)
{
	if (!vars->encoder) return FALSE;

	Int resultReceive = avcodec_receive_packet(vars->encoder, vars->packet);
	if (resultReceive < 0) return FALSE;

	packetData = Binary(av_packet_get_data(vars->packet), av_packet_get_size(vars->packet));
	av_packet_unref(vars->packet);

	if (vars->operationWait) Flag(OPERATION_WAIT_TIME).waitSet();

	Array<Double> dequeued = vars->times.dequeue(1);
	Bool ok = dequeued.size() == 1;
	if (ok) time = dequeued[0];
	return ok;
}
